#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <poll.h>
#include <fcntl.h>

#define MAX_CLIENTS 10
#define SERVER_INDEX 0
#define CLIENT_BEGIN_INDEX 1
#define BUFFER_SIZE 1024
#define PORT 8080

// 函数用于处理客户端消息
void handle_client(int sockfd, struct sockaddr_in client_addr, socklen_t client_len) {
    char buffer[BUFFER_SIZE];
    ssize_t bytes_received;
    const char *response = "ACK from Server";

    while (1) {
        // 清零缓冲区
        memset(buffer, 0, BUFFER_SIZE);

        // 接收数据
        bytes_received = recvfrom(sockfd, buffer, BUFFER_SIZE, 0, (struct sockaddr *)&client_addr, &client_len);
        if (bytes_received < 0) {
            perror("recvfrom");
            break;
        } else if (bytes_received == 0) {
            printf("Client disconnected\n");
            break;
        }

        printf("Received message from client: %s\n", buffer);

        // 发送响应
        sendto(sockfd, response, strlen(response), 0, (struct sockaddr *)&client_addr, client_len);
    }
}

int main() {
    int sockfd, new_fd;
    struct sockaddr_in serv_addr, client_addr;
    socklen_t client_len = sizeof(client_addr);
    struct pollfd fds[MAX_CLIENTS + 1]; // 多一个用于监听套接字
    int max_fds = 1; // 初始化已知的最大文件描述符数量
    int i, num_events;

    // 创建UDP套接字
    if ((sockfd = socket(AF_INET, SOCK_DGRAM, 0)) < 0) {
        perror("Socket creation failed");
        exit(EXIT_FAILURE);
    }

    // 设置为非阻塞模式，这对于使用poll很重要
    if (fcntl(sockfd, F_SETFL, O_NONBLOCK) == -1) {
        perror("Setting non-blocking mode failed");
        exit(EXIT_FAILURE);
    }

    memset(&serv_addr, 0, sizeof(serv_addr));
    serv_addr.sin_family = AF_INET;
    serv_addr.sin_addr.s_addr = INADDR_ANY;
    serv_addr.sin_port = htons(PORT);

    // 绑定套接字
    if (bind(sockfd, (struct sockaddr *)&serv_addr, sizeof(serv_addr)) < 0) {
        perror("Socket bind failed");
        exit(EXIT_FAILURE);
    }

    // 初始化pollfd数组
    fds[SERVER_INDEX].fd = sockfd;
    fds[SERVER_INDEX].events = POLLIN; // 监听套接字上是否有数据到达
    int index = CLIENT_BEGIN_INDEX;
    while (1) {
        num_events = poll(fds, max_fds, -1); // -1表示无限期等待

        if (num_events < 0) {
            perror("poll");
            exit(EXIT_FAILURE);
        }

        for (i = 0; i < max_fds; i++) {
            if (fds[i].revents & POLLIN) {
                if (i == SERVER_INDEX) { // 监听套接字有活动
                    // 处理新连接
                    fds[index].fd = sockfd;
                    fds[index].events = POLLIN; // 监听套接字上是否有数据到达
                    ++index;
                } else { // 客户端套接字有活动
                    // 实际应用中，需要根据fds[i].fd找到对应的客户端处理
                    // 此处简化处理，直接调用处理函数，实际应用需维护客户端状态
                    handle_client(fds[i].fd, client_addr, client_len);
                }
            }
        }
        max_fds = index;
    }

    // 实际应用中应关闭所有打开的套接字
    close(sockfd);

    return 0;
}